#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Betaflight黑盒日志分析器
用于解析.bbl文件并绘制PID曲线
"""

import os
import sys
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.font_manager as fm
from typing import Dict, List, Tuple, Optional

try:
    from orangebox import Parser
except ImportError:
    print("错误：未找到orangebox库，请安装：pip install orangebox")
    sys.exit(1)

class BetaflightAnalyzer:
    """Betaflight黑盒日志分析器"""

    def __init__(self, log_file: str):
        """初始化分析器

        Args:
            log_file: .bbl日志文件路径
        """
        self.log_file = log_file
        self.parser = None
        self.data = None
        self.field_names = []

        # 设置中文字体
        plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans']
        plt.rcParams['axes.unicode_minus'] = False

    def load_log(self, log_index: int = 1) -> bool:
        """加载日志文件

        Args:
            log_index: 日志索引（如果文件包含多个日志）

        Returns:
            bool: 加载是否成功
        """
        try:
            print(f"正在加载日志文件: {self.log_file}")
            self.parser = Parser.load(self.log_file, log_index)
            self.field_names = self.parser.field_names

            print(f"日志数量: {self.parser.reader.log_count}")
            print(f"字段数量: {len(self.field_names)}")
            print(f"可用字段: {', '.join(self.field_names[:10])}{'...' if len(self.field_names) > 10 else ''}")

            return True

        except Exception as e:
            print(f"加载日志文件失败: {e}")
            return False

    def extract_data(self) -> bool:
        """提取所有数据

        Returns:
            bool: 提取是否成功
        """
        try:
            print("正在提取数据...")
            self.data = {}

            # 初始化数据字典
            for field_name in self.field_names:
                self.data[field_name] = []

            # 提取所有数据 - 使用正确的orangebox API
            frame_count = 0
            for frame in self.parser.frames():
                frame_count += 1
                # frame.data是一个tuple，需要按索引访问
                for i, field_name in enumerate(self.field_names):
                    if i < len(frame.data):
                        self.data[field_name].append(frame.data[i])
                    else:
                        self.data[field_name].append(None)

            # 转换为numpy数组，处理None值
            for field_name in self.data:
                # 过滤掉None值
                valid_data = [x for x in self.data[field_name] if x is not None]
                if valid_data:
                    try:
                        self.data[field_name] = np.array(valid_data)
                    except:
                        # 如果无法转换为数值数组，保持原始列表
                        self.data[field_name] = valid_data
                else:
                    self.data[field_name] = np.array([])

            print(f"数据提取完成，共 {frame_count} 帧，{len(self.data)} 个字段")
            return True

        except Exception as e:
            print(f"数据提取失败: {e}")
            import traceback
            traceback.print_exc()
            return False

    def get_available_fields(self) -> List[str]:
        """获取可用的字段列表"""
        if self.data:
            return list(self.data.keys())
        return []

    def plot_pid_curves(self, output_dir: str = ".") -> bool:
        """绘制PID曲线

        Args:
            output_dir: 输出目录

        Returns:
            bool: 绘制是否成功
        """
        if not self.data:
            print("错误：没有可用数据")
            return False

        try:
            # 查找时间字段
            time_field = None
            for field in ['time', 'loopIteration', 'iteration']:
                if field in self.data:
                    time_field = field
                    break

            if time_field is None:
                # 如果没有时间字段，创建一个基于索引的时间轴
                time_data = np.arange(len(next(iter(self.data.values()))))
                time_label = "数据点索引"
            else:
                time_data = self.data[time_field]
                time_label = "时间 (μs)" if time_field == 'time' else "循环迭代"

            # 查找PID相关字段
            pid_fields = self._find_pid_fields()

            if not pid_fields:
                print("警告：未找到PID相关字段")
                self._plot_all_available_data(time_data, time_label, output_dir)
                return True

            # 绘制PID曲线
            self._plot_pid_data(time_data, time_label, pid_fields, output_dir)

            return True

        except Exception as e:
            print(f"绘制PID曲线失败: {e}")
            return False

    def _find_pid_fields(self) -> Dict[str, List[str]]:
        """查找PID相关字段"""
        pid_fields = {
            'setpoint': [],
            'gyro': [],
            'pid_output': [],
            'pid_p': [],
            'pid_i': [],
            'pid_d': []
        }

        # 常见的PID字段模式
        patterns = {
            'setpoint': ['setpoint', 'rcCommand'],
            'gyro': ['gyroADC', 'gyro'],
            'pid_output': ['motor', 'pidSum'],
            'pid_p': ['pidP'],
            'pid_i': ['pidI'],
            'pid_d': ['pidD']
        }

        for category, pattern_list in patterns.items():
            for field_name in self.data.keys():
                for pattern in pattern_list:
                    if pattern.lower() in field_name.lower():
                        pid_fields[category].append(field_name)

        return pid_fields

    def _plot_pid_data(self, time_data: np.ndarray, time_label: str,
                       pid_fields: Dict[str, List[str]], output_dir: str):
        """绘制PID数据"""
        # 创建子图
        fig, axes = plt.subplots(3, 2, figsize=(15, 12))
        fig.suptitle('Betaflight PID 分析图表', fontsize=16, fontweight='bold')

        # 轴标签
        axis_labels = ['Roll', 'Pitch', 'Yaw']

        # 绘制设定点 vs 实际值
        for i, axis in enumerate(axis_labels):
            ax = axes[i, 0]

            # 查找对应轴的字段
            setpoint_fields = [f for f in pid_fields['setpoint'] if str(i) in f or axis.lower() in f.lower()]
            gyro_fields = [f for f in pid_fields['gyro'] if str(i) in f or axis.lower() in f.lower()]

            if setpoint_fields:
                ax.plot(time_data, self.data[setpoint_fields[0]],
                       label=f'{axis} 设定点', linewidth=1, alpha=0.8)

            if gyro_fields:
                ax.plot(time_data, self.data[gyro_fields[0]],
                       label=f'{axis} 实际值', linewidth=1, alpha=0.8)

            ax.set_title(f'{axis}轴 - 设定点 vs 实际值')
            ax.set_xlabel(time_label)
            ax.set_ylabel('角速度 (°/s)')
            ax.legend()
            ax.grid(True, alpha=0.3)

        # 绘制PID输出
        for i, axis in enumerate(axis_labels):
            ax = axes[i, 1]

            # 查找PID输出字段
            p_fields = [f for f in pid_fields['pid_p'] if str(i) in f or axis.lower() in f.lower()]
            i_fields = [f for f in pid_fields['pid_i'] if str(i) in f or axis.lower() in f.lower()]
            d_fields = [f for f in pid_fields['pid_d'] if str(i) in f or axis.lower() in f.lower()]

            if p_fields:
                ax.plot(time_data, self.data[p_fields[0]],
                       label='P项', linewidth=1, alpha=0.8)

            if i_fields:
                ax.plot(time_data, self.data[i_fields[0]],
                       label='I项', linewidth=1, alpha=0.8)

            if d_fields:
                ax.plot(time_data, self.data[d_fields[0]],
                       label='D项', linewidth=1, alpha=0.8)

            ax.set_title(f'{axis}轴 - PID输出')
            ax.set_xlabel(time_label)
            ax.set_ylabel('PID输出')
            ax.legend()
            ax.grid(True, alpha=0.3)

        plt.tight_layout()

        # 保存图表
        output_file = os.path.join(output_dir, f"pid_analysis_{os.path.basename(self.log_file)}.png")
        plt.savefig(output_file, dpi=300, bbox_inches='tight')
        print(f"PID分析图表已保存: {output_file}")

        plt.show()

    def _plot_all_available_data(self, time_data: np.ndarray, time_label: str, output_dir: str):
        """绘制所有可用数据（当找不到PID字段时）"""
        print("绘制所有可用数据字段...")

        # 选择前12个数值字段进行绘制
        numeric_fields = []
        for field_name, field_data in self.data.items():
            if len(field_data) > 0 and np.issubdtype(field_data.dtype, np.number):
                numeric_fields.append(field_name)
            if len(numeric_fields) >= 12:
                break

        if not numeric_fields:
            print("错误：未找到可绘制的数值字段")
            return

        # 创建子图
        rows = min(4, len(numeric_fields))
        cols = min(3, (len(numeric_fields) + rows - 1) // rows)

        fig, axes = plt.subplots(rows, cols, figsize=(15, 12))
        fig.suptitle('Betaflight 黑盒数据分析', fontsize=16, fontweight='bold')

        if rows == 1 and cols == 1:
            axes = [axes]
        elif rows == 1 or cols == 1:
            axes = axes.flatten()
        else:
            axes = axes.flatten()

        for i, field_name in enumerate(numeric_fields[:len(axes)]):
            ax = axes[i]
            ax.plot(time_data, self.data[field_name], linewidth=1, alpha=0.8)
            ax.set_title(field_name)
            ax.set_xlabel(time_label)
            ax.grid(True, alpha=0.3)

        # 隐藏多余的子图
        for i in range(len(numeric_fields), len(axes)):
            axes[i].set_visible(False)

        plt.tight_layout()

        # 保存图表
        output_file = os.path.join(output_dir, f"blackbox_data_{os.path.basename(self.log_file)}.png")
        plt.savefig(output_file, dpi=300, bbox_inches='tight')
        print(f"数据分析图表已保存: {output_file}")

        plt.show()

def main():
    """主函数"""
    # 日志文件路径
    log_files = [
        "d:/Git/FlightRecord/SFPV_85/record/btfl_001.bbl",
        "d:/Git/FlightRecord/SFPV_85/record/btfl_002.bbl",
        "d:/Git/FlightRecord/SFPV_85/record/btfl_all.bbl"
    ]

    # 输出目录
    output_dir = "d:/Git/FlightRecord/SFPV_85/analysis_results"
    os.makedirs(output_dir, exist_ok=True)

    for log_file in log_files:
        if not os.path.exists(log_file):
            print(f"警告：文件不存在 {log_file}")
            continue

        print(f"\n{'='*60}")
        print(f"分析文件: {os.path.basename(log_file)}")
        print(f"{'='*60}")

        # 创建分析器
        analyzer = BetaflightAnalyzer(log_file)

        # 加载日志
        if not analyzer.load_log():
            continue

        # 提取数据
        if not analyzer.extract_data():
            continue

        # 显示可用字段
        fields = analyzer.get_available_fields()
        print(f"\n可用字段 ({len(fields)} 个):")
        for i, field in enumerate(fields[:20]):  # 只显示前20个
            print(f"  {i+1:2d}. {field}")
        if len(fields) > 20:
            print(f"  ... 还有 {len(fields)-20} 个字段")

        # 绘制PID曲线
        analyzer.plot_pid_curves(output_dir)

        print(f"\n文件 {os.path.basename(log_file)} 分析完成")

    print(f"\n所有分析结果已保存到: {output_dir}")

if __name__ == "__main__":
    main()